#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from retrying import retry
from tqdm import tqdm
import requests
import json


def search_pubmed_count(term: str):
    url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi'
    payload = {"db": "pubmed", "term": term, "retmode": "json"}
    r = requests.post(url, stream=True, data=payload)
    if r.status_code == 200:
        result = json.loads(r.text)
        pmid_count = int(result["esearchresult"]["count"])
        return pmid_count
    else:
        print("Web Error. Retrying ...")
        raise IOError("Web Error.")


@retry
def get_pmids(term: str, retstart: str, retmax: str):
    url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi'
    payload = {"db": "pubmed", "term": term, "retmax": retmax, "retstart": retstart, "retmode": "json"}
    r = requests.post(url, stream=True, data=payload)
    if r.status_code == 200:
        result = json.loads(r.text)
        return result["esearchresult"]["idlist"]
    else:
        print("Web Error. Retrying ...")
        raise IOError("Web Error.")


def search_pubmed(term: str, retmax: int, outpath: str):
    outfile = open(outpath, "w")
    pmid_count = search_pubmed_count(term)
    for i in tqdm(range(0, pmid_count, retmax)):
        tmp_pmids = get_pmids(term, str(i), str(retmax))
        for item in tmp_pmids:
            outfile.write(item + "\n")
            outfile.flush()
    outfile.close()


if __name__ == '__main__':
    search_term='''"Lung Neoplasms"[MeSH Terms] OR "Adenocarcinoma of Lung"[MeSH Terms] OR "Bronchial Neoplasms"[MeSH Terms] OR "Carcinoma, Bronchogenic"[MeSH Terms] OR "Lung Neoplasms"[All Fields] OR "Adenocarcinoma of Lung"[All Fields] OR "Bronchial Neoplasms"[All Fields] OR "Carcinoma, Bronchogenic"[All Fields] OR "Primary malignant neoplasm of lung"[All Fields] OR "Large cell carcinoma of lung"[All Fields] OR "Non-Hodgkin's lymphoma of lung"[All Fields] OR "Non-small cell lung cancer"[All Fields] OR "Pleuropulmonary blastoma"[All Fields] OR "Primary adenocarcinoma of lung"[All Fields] OR "Primary malignant neoplasm of left lung"[All Fields] OR "Primary malignant neoplasm of lower lobe of left lung"[All Fields] OR "Primary malignant neoplasm of lower lobe, bronchus or lung"[All Fields] OR "Primary malignant neoplasm of right lung"[All Fields] OR "Small cell carcinoma of lung"[All Fields] OR "Squamous cell carcinoma of lung"[All Fields] OR "Malignant neoplasms of bronchus or lung"[All Fields] OR "Malignant neoplasm of bronchus and lung"[All Fields] OR "Mesothelioma, Malignant"[All Fields] OR "Multiple Pulmonary Nodules"[All Fields] OR "Pancoast Syndrome"[All Fields] OR "Pulmonary Blastoma"[All Fields] OR "Pulmonary Sclerosing Hemangioma"[All Fields] OR "Adenocarcinoma, Bronchiolo-Alveolar"[All Fields] OR "Carcinoma, Non-Small-Cell Lung"[All Fields] OR "Small Cell Lung Carcinoma"[All Fields] OR "Acinar cell cystadenocarcinoma of lung"[All Fields] OR "Alpha heavy chain disease, respiratory form"[All Fields] OR "Epithelioid hemangioendothelioma of lung"[All Fields] OR "Pancoast tumor"[All Fields] OR "Primary acinar cell carcinoma of lung"[All Fields] OR "Primary carcinosarcoma of lung"[All Fields] OR "Primary malignant neoplasm of hilus of lung"[All Fields] OR "Primary malignant neuroendocrine neoplasm of lung"[All Fields] OR "Primary pleomorphic carcinoma of lung"[All Fields] OR "Primary pseudosarcomatous carcinoma of lung"[All Fields] OR "Primary signet ring cell carcinoma of lung"[All Fields] OR "Primary spindle cell carcinoma of lung"[All Fields] OR "Primary undifferentiated carcinoma of lung"[All Fields] OR "Giant cell carcinoma of lung"[All Fields] OR "Primary large cell carcinoma of lower lobe of left lung"[All Fields] OR "Primary large cell carcinoma of lower lobe of right lung"[All Fields] OR "Primary large cell carcinoma of upper lobe of left lung"[All Fields] OR "Primary large cell carcinoma of upper lobe of right lung"[All Fields] OR "Diffuse non-Hodgkin's lymphoma of lung"[All Fields] OR "Follicular non-Hodgkin's lymphoma of lung"[All Fields] OR "Anaplastic lymphoma kinase fusion oncogene negative non-small cell lung cancer"[All Fields] OR "Anaplastic lymphoma kinase fusion oncogene positive non-small cell lung cancer"[All Fields] OR "Non-small cell lung cancer with mutation in epidermal growth factor receptor"[All Fields] OR "Non-small cell lung cancer without mutation in epidermal growth factor receptor"[All Fields] OR "Non-small cell lung cancer, negative for epidermal growth factor receptor expression"[All Fields] OR "Non-small cell lung cancer, positive for epidermal growth factor receptor expression"[All Fields] OR "Non-small cell lung carcinoma with neuregulin 1 gene fusion"[All Fields] OR "Nonsquamous nonsmall cell neoplasm of lung"[All Fields] OR "Primary non-small cell carcinoma of left lung"[All Fields] OR "Primary non-small cell carcinoma of right lung"[All Fields] OR "Reactive oxygen species 1 negative non-small cell lung cancer"[All Fields] OR "Reactive oxygen species 1 positive non-small cell lung cancer"[All Fields] OR "Squamous non-small cell lung cancer"[All Fields] OR "Pleuropulmonary blastoma type I"[All Fields] OR "Pleuropulmonary blastoma type II"[All Fields] OR "Pleuropulmonary blastoma type III"[All Fields] OR "Primary adenocarcinoma of lower lobe of left lung"[All Fields] OR "Primary adenocarcinoma of lower lobe of right lung"[All Fields] OR "Primary adenocarcinoma of middle lobe of right lung"[All Fields] OR "Primary adenocarcinoma of upper lobe of left lung"[All Fields] OR "Primary adenocarcinoma of upper lobe of right lung"[All Fields] OR "Primary adenosquamous carcinoma of lung"[All Fields] OR "Primary clear cell adenocarcinoma of lung"[All Fields] OR "Primary fetal adenocarcinoma of lung"[All Fields] OR "Primary mixed subtype adenocarcinoma of lung"[All Fields] OR "Primary mucinous adenocarcinoma of lung"[All Fields] OR "Primary mucinous bronchiolo-alveolar carcinoma of lung"[All Fields] OR "Primary mucinous cystadenocarcinoma of lung"[All Fields] OR "Primary non-mucinous bronchiolo-alveolar carcinoma of lung"[All Fields] OR "Primary papillary adenocarcinoma of lung"[All Fields] OR "Primary salivary gland type carcinoma of lung"[All Fields] OR "Primary solid carcinoma of lung"[All Fields] OR "Primary malignant neoplasm of bilateral lungs"[All Fields] OR "Primary malignant neoplasm of upper lobe of left lung"[All Fields] OR "Primary small cell carcinoma of left lung"[All Fields] OR "Squamous cell carcinoma of left lung"[All Fields] OR "Primary malignant neoplasm of bronchus of left lower lobe"[All Fields] OR "Primary non-small cell carcinoma of lower lobe of left lung"[All Fields] OR "Primary small cell carcinoma of lower lobe of left lung"[All Fields] OR "Primary squamous cell carcinoma of lower lobe of left lung"[All Fields] OR "Carcinoma of lower lobe, bronchus or lung"[All Fields] OR "Primary malignant neoplasm of right lower lobe of lung"[All Fields] OR "Primary malignant neoplasm of right middle lobe of lung"[All Fields] OR "Primary malignant neoplasm of upper lobe of right lung"[All Fields] OR "Primary small cell carcinoma of right lung"[All Fields] OR "Squamous cell carcinoma of right lung"[All Fields] OR "Extensive stage primary small cell carcinoma of lung"[All Fields] OR "Oat cell carcinoma of lung"[All Fields] OR "Primary basaloid squamous cell carcinoma of lung"[All Fields] OR "Primary clear cell squamous cell carcinoma of lung"[All Fields] OR "Primary papillary squamous cell carcinoma of lung"[All Fields] OR "Primary small cell non-keratinizing squamous cell carcinoma of lung"[All Fields] OR "Adenocarcinoma of bronchus or lung"[All Fields] OR "Small cell carcinoma of bronchus or lung"[All Fields] OR "Squamous cell carcinoma of bronchus or lung"[All Fields] OR "Large cell carcinoma of bronchus or lung"[All Fields] OR "Carcinoid or other malignant neuroendocrine neoplasms of bronchus or lung"[All Fields] OR "Unspecified malignant epithelial neoplasm of bronchus or lung"[All Fields] OR "Other specified malignant neoplasms of bronchus or lung"[All Fields] OR "Malignant neoplasms of bronchus or lung, unspecified"[All Fields] OR "Upper lobe, bronchus or lung"[All Fields] OR "Middle lobe, bronchus or lung"[All Fields] OR "Lower lobe, bronchus or lung"[All Fields] OR "Overlapping lesion of bronchus and lung"[All Fields] OR "Bronchus or lung, unspecified"[All Fields] OR "Pulmonary Neoplasms"[All Fields] OR "Neoplasms, Lung"[All Fields] OR "Lung Neoplasm"[All Fields] OR "Neoplasm, Lung"[All Fields] OR "Neoplasms, Pulmonary"[All Fields] OR "Neoplasm, Pulmonary"[All Fields] OR "Pulmonary Neoplasm"[All Fields] OR "Lung Cancer"[All Fields] OR "Cancer, Lung"[All Fields] OR "Cancers, Lung"[All Fields] OR "Lung Cancers"[All Fields] OR "Pulmonary Cancer"[All Fields] OR "Cancer, Pulmonary"[All Fields] OR "Cancers, Pulmonary"[All Fields] OR "Pulmonary Cancers"[All Fields] OR "Cancer of the Lung"[All Fields] OR "Cancer of Lung"[All Fields] OR "Lung Adenocarcinomas"[All Fields] OR "Lung Adenocarcinoma"[All Fields] OR "Adenocarcinoma, Lung"[All Fields] OR "Adenocarcinomas, Lung"[All Fields] OR "Neoplasms, Bronchial"[All Fields] OR "Bronchial Neoplasm"[All Fields] OR "Neoplasm, Bronchial"[All Fields] OR "Bronchogenic Carcinoma"[All Fields] OR "Bronchogenic Carcinomas"[All Fields] OR "Carcinomas, Bronchogenic"[All Fields] OR "Carcinoma, Bronchial"[All Fields] OR "Bronchial Carcinoma"[All Fields] OR "Bronchial Carcinomas"[All Fields] OR "Carcinomas, Bronchial"[All Fields] OR "Non-Hodgkin lymphoma of lung"[All Fields] OR "Pleuro-pulmonary blastoma"[All Fields] OR "Primary malignant neoplasm of left lower lobe of lung"[All Fields] OR "Small cell lung cancer"[All Fields]'''
    # search_pubmed(search_term)
    search_pubmed(search_term, 100000, "/Users/genesis/Downloads/test.txt")
